package ymratelimit

import (
	"encoding/json"
	"sync/atomic"
	"time"
	"unsafe"
)

// Limiter /*
type Limiter interface {
	TakeAvailableWithNow(now int64) bool
	TakeAvailable() bool
	GetCapacity() int64
	GetLegacyCapacity() int64
}

/*
*

	token bucket algorithm
*/
type tokenBucket struct {
	capacity        int64          // bucket capacity
	fillInterval    int64          // fill tokens interval
	tokenBucketStat unsafe.Pointer // current status
}

/*
*

	current status
*/
type tokenBucketStat struct {
	nextTokenTimestamp int64
	keepCapacity       int64 // token remains this time slot
}

// TakeAvailableWithNow /*
func (t *tokenBucket) TakeAvailableWithNow(now int64) bool {

	// for to ensure LoadPointer and CompareAndSwapPointer are atomic
	taken := false

	for !taken {
		lastTokenBucketStatPointer := atomic.LoadPointer(&t.tokenBucketStat)
		lastTokenBucketStat := (*tokenBucketStat)(lastTokenBucketStatPointer)

		if now > lastTokenBucketStat.nextTokenTimestamp {
			newStat := tokenBucketStat{}
			newStat.nextTokenTimestamp = lastTokenBucketStat.nextTokenTimestamp + t.fillInterval*(((now-lastTokenBucketStat.nextTokenTimestamp)/t.fillInterval)+1)
			newStat.keepCapacity = t.capacity - 1
			taken = atomic.CompareAndSwapPointer(&t.tokenBucketStat, lastTokenBucketStatPointer, unsafe.Pointer(&newStat))
		} else {
			// tokens remains, just consume.
			if lastTokenBucketStat.keepCapacity > 0 {
				newStat := tokenBucketStat{}
				newStat.nextTokenTimestamp = lastTokenBucketStat.nextTokenTimestamp
				newStat.keepCapacity = lastTokenBucketStat.keepCapacity - 1
				taken = atomic.CompareAndSwapPointer(&t.tokenBucketStat, lastTokenBucketStatPointer, unsafe.Pointer(&newStat))
			} else {
				// no tokens any more
				break
			}
		}
	}

	return taken
}

func (t *tokenBucket) TakeAvailable() bool {
	return t.TakeAvailableWithNow(time.Now().UnixNano())
}

func (t *tokenBucket) GetCapacity() int64 {
	return t.capacity
}

func (t *tokenBucket) GetLegacyCapacity() int64 {
	lastTokenBucketStatPointer := atomic.LoadPointer(&t.tokenBucketStat)
	lastTokenBucketStat := (*tokenBucketStat)(lastTokenBucketStatPointer)
	return lastTokenBucketStat.keepCapacity
}

func (t *tokenBucket) MarshalJSON() ([]byte, error) {
	object := map[string]interface{}{}
	object["capacity"] = t.capacity
	lastTokenBucketStatPointer := atomic.LoadPointer(&t.tokenBucketStat)
	lastTokenBucketStat := (*tokenBucketStat)(lastTokenBucketStatPointer)
	object["keepCapacity"] = lastTokenBucketStat.keepCapacity
	return json.Marshal(object)
}

// NewTokenBucket /*
func NewTokenBucket(fillInterval time.Duration, capacity int64) Limiter {
	fillIntervalInt := int64(fillInterval)
	l := &tokenBucket{
		fillInterval: fillIntervalInt,
		capacity:     capacity,
	}
	tokenBucketStat := tokenBucketStat{
		nextTokenTimestamp: time.Now().UnixNano(),
		keepCapacity:       capacity,
	}
	l.tokenBucketStat = unsafe.Pointer(&tokenBucketStat)
	return l
}
